//
//    node-l3qd - The L3q daemon that runs on each compute node.
//    Copyright (C) 2022-2023  Marcus Pedersén marcus@marcux.org
//
//    This program is free software: you can redistribute it and/or modify
//    it under the terms of the GNU General Public License as published by
//    the Free Software Foundation, either version 3 of the License, or
//    (at your option) any later version.
//
//    This program is distributed in the hope that it will be useful,
//    but WITHOUT ANY WARRANTY; without even the implied warranty of
//    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
//    GNU General Public License for more details.
//
//    You should have received a copy of the GNU General Public License
//    along with this program.  If not, see <http://www.gnu.org/licenses/>.
//

use actix_web::{error, web, Error};
use rusqlite::{Statement, params};
use crate::comm::node::TaskStatus;

pub type Pool = r2d2::Pool<r2d2_sqlite::SqliteConnectionManager>;
pub type Connection = r2d2::PooledConnection<r2d2_sqlite::SqliteConnectionManager>;
pub type DbResponseResult = Result<Vec<DbResult>, rusqlite::Error>;

#[derive(Debug)]
pub enum DbResult {
    InsertSuccess,
    DeleteSuccess,
    Validate {
        ip: String,
        name: String,
        key: String,
    },
    Task {
        id: i64,
        jobid: u64,
        taskid: u64,
        unitname: String,
        start_time: String,
        end_time: String,
        mono_start_time: u64,
        mono_end_time: u64,
        active_state: String,
        sub_state: String,
        exit_code: u16,
        main_pid: u32,
        memory_peak: u64,
        pids_peak: u32,
        stdout: String,
        stderr: String,
     },
}

impl DbResult {
    pub fn to_task_status(task: &DbResult) -> TaskStatus {
        match task {
            DbResult::Task {
                id: _,
                jobid,
                taskid,
                unitname,
                start_time,
                end_time,
                mono_start_time,
                mono_end_time,
                active_state,
                sub_state,
                exit_code,
                main_pid,
                memory_peak,
                pids_peak,
                stdout,
                stderr,
            } => {
                TaskStatus {
                    jobid: *jobid,
                    taskid: *taskid,
                    unitname: unitname.to_string().clone(),
                    start_time: start_time.to_string().clone(),
                    end_time: end_time.to_string().clone(),
                    mono_start_time: *mono_start_time,
                    mono_end_time: *mono_end_time,
                    active_state: active_state.to_string().clone(),
                    sub_state: sub_state.to_string().clone(),
                    exit_code: *exit_code,
                    main_pid: *main_pid,
                    memory_peak: *memory_peak,
                    pids_peak: *pids_peak,
                    stdout: stdout.to_string().clone(),
                    stderr: stderr.to_string().clone(),
                }
            },
            _ => TaskStatus::new(),
        }
    }
}

/// Queries executed by this
/// module
pub enum DbQuery {
    GetL3qdValidate,
    GetNodeL3qdValidate,
    InsertValidate {
        ip: String,
        name: String,
        key: String,
    },
    InsertTask {
        jobid: u64,
        taskid: u64,
        unitname: String,
        start_time: String,
        end_time: String,
        mono_start_time: u64,
        mono_end_time: u64,
        active_state: String,
        sub_state: String,
        exit_code: u16,
        main_pid: u32,
        memory_peak: u64,
        pids_peak: u32,
        stdout: String,
        stderr: String,
    },
    DeleteAllTasks,
    GetAllTasks,
}

impl DbQuery {
    pub fn to_insert_task(t: &TaskStatus) -> DbQuery {
        DbQuery::InsertTask {
            jobid: t.jobid,
            taskid: t.taskid,
            unitname: t.unitname.clone(),
            start_time: t.start_time.clone(),
            end_time: t.end_time.clone(),
            mono_start_time: t.mono_start_time,
            mono_end_time: t.mono_end_time,
            active_state: t.active_state.clone(),
            sub_state: t.sub_state.clone(),
            exit_code: t.exit_code,
            main_pid: t.main_pid,
            memory_peak: t.memory_peak,
            pids_peak: t.pids_peak,
            stdout: t.stdout.clone(),
            stderr: t.stderr.clone(),
        }
    }
}


/// To be used from validate process
/// that is a sync function.
fn get_db_connection(db_path: &String) -> Result<rusqlite::Connection, rusqlite::Error> {
    Ok(rusqlite::Connection::open(db_path)?)
}


/// To be used from validate process
/// that is a sync function.
/// Inserts a new validation key
/// If it exists it will be over written.
pub fn insert_validate(db_path: &String, validate: DbQuery) -> DbResponseResult {
    let conn = get_db_connection(db_path)?;
    
    match validate {
        DbQuery::InsertValidate{ip, name, key} => {
            let mut stmt = conn.prepare(
                "SELECT ip, name, key
                 FROM validate
                 WHERE name = ?1",
            )?;

            let exists: Vec::<DbResult> = stmt.query_map(params![name], |row| {
                Ok(DbResult::Validate {
                    ip: row.get(0)?,
                    name: row.get(1)?,
                    key: row.get(2)?,
                })
            })
                .and_then(Iterator::collect)?;
            
            let mut stmt;
            if exists.len() >= 1 {
                stmt = conn.prepare(
                    "UPDATE validate
                     SET ip = ?1,
                         key = ?2
                     WHERE name = ?3")?;
                stmt.execute(params![ip, key, name])?;
            }
            else {
                stmt = conn.prepare(
                    "INSERT INTO validate
                     (ip, name, key)
                     VALUES(?1, ?2, ?3)")?;
                stmt.execute(params![ip, name, key])?;
            }




            Ok(vec![DbResult::InsertSuccess])
        },
        _ => Err(rusqlite::Error::InvalidQuery),
    }
}



/// Returns the validate object
/// for the l3q-daemon
fn get_l3qd_validate(conn: Connection) -> DbResponseResult {
    let stmt = conn.prepare(
        "SELECT ip, name, key
         FROM validate
         WHERE name = 'l3qd'",
    )?;

    get_validate_result(stmt)
}


/// Returns this host (node-l3qd)
/// validation object.
fn get_node_l3qd_validate(conn: Connection) -> DbResponseResult {
    let stmt = conn.prepare(
        "SELECT ip, name, key
         FROM validate
         WHERE name = 'self'",
    )?;

    get_validate_result(stmt)
}


/// Interpretates validate results
/// from the database.
fn get_validate_result(mut stmt: Statement) -> DbResponseResult {
    let res: Vec::<DbResult> = stmt.query_map([], |row| {
        Ok(DbResult::Validate {
            ip: row.get(0)?,
            name: row.get(1)?,
            key: row.get(2)?,
        })
    })
        .and_then(Iterator::collect)?;

    Ok(res)
}


/// Returns all tasks stored in
/// in database as task objects.
fn get_all_tasks(conn: Connection) -> DbResponseResult {
    let mut stmt = conn.prepare(
        "SELECT * FROM status_task",
    )?;

    let res: Vec::<DbResult> = stmt.query_map([], |row| {
        Ok(DbResult::Task {
            id: row.get(0)?,
            jobid: row.get(1)?,
            taskid: row.get(2)?,
            unitname: row.get(3)?,
            start_time: row.get(4)?,
            end_time: row.get(5)?,
            mono_start_time: row.get(6)?,
            mono_end_time: row.get(7)?,
            active_state: row.get(8)?,
            sub_state: row.get(9)?,
            exit_code: row.get(10)?,
            main_pid: row.get(11)?,
            memory_peak: row.get(12)?,
            pids_peak: row.get(13)?,
            stdout: row.get(14)?,
            stderr: row.get(15)?,
        })
    })
        .and_then(Iterator::collect)?;

    Ok(res)
}


/// Deletes all tasks
/// from database when
/// central daemon is online
/// and has received all tasks
fn delete_all_tasks(conn: Connection) -> DbResponseResult {
    let mut stmt = conn.prepare("DELETE FROM status_task")?;

    stmt.execute(params![])?;

    Ok(vec![DbResult::DeleteSuccess])
}


/// Inserts a new task or updates
/// task if already exists
fn insert_task(conn: Connection, task: DbQuery) -> DbResponseResult {
    match task {
        DbQuery::InsertTask {
            jobid,
            taskid,
            unitname,
            start_time,
            end_time,
            mono_start_time,
            mono_end_time,
            active_state,
            sub_state,
            exit_code,
            main_pid,
            memory_peak,
            pids_peak,
            stdout,
            stderr,
        } => {
            let mut stmt = conn.prepare(
                "SELECT *
                 FROM status_task
                 WHERE jobid = ?1 AND taskid = ?2",
            )?;

            let exists: Vec::<DbResult> = stmt.query_map(params![jobid, taskid], |row| {
                Ok(DbResult::Task {
                    id: row.get(0)?,
                    jobid: row.get(1)?,
                    taskid: row.get(2)?,
                    unitname: row.get(3)?,
                    start_time: row.get(4)?,
                    end_time: row.get(5)?,
                    mono_start_time: row.get(6)?,
                    mono_end_time: row.get(7)?,
                    active_state: row.get(8)?,
                    sub_state: row.get(9)?,
                    exit_code: row.get(10)?,
                    main_pid: row.get(11)?,
                    memory_peak: row.get(12)?,
                    pids_peak: row.get(13)?,
                    stdout: row.get(14)?,
                    stderr: row.get(15)?,
                })
            })
                .and_then(Iterator::collect)?;
            
            let mut stmt;
            if exists.len() >= 1 {
                if sub_state == String::from("failed") || sub_state == String::from("exited") {
                    stmt = conn.prepare(
                        "UPDATE status_task
                         SET unitname = ?1,
                             start_time = ?2,
                             end_time = ?3,
                             mono_start_time = ?4,
                             mono_end_time = ?5,
                             active_state = ?6,
                             sub_state = ?7,
                             exit_code = ?8,
                             stdout = ?9,
                             stderr = ?10
                         WHERE jobid = ?11 AND taskid = ?12")?;
                    stmt.execute(params![
                        unitname,
                        start_time,
                        end_time,
                        mono_start_time,
                        mono_end_time,
                        active_state,
                        sub_state,
                        exit_code,
                        stdout,
                        stderr,
                        jobid,
                        taskid
                    ])?;
                }
                else {
                    stmt = conn.prepare(
                        "UPDATE status_task
                         SET unitname = ?1,
                             start_time = ?2,
                             end_time = ?3,
                             mono_start_time = ?4,
                             mono_end_time = ?5,
                             active_state = ?6,
                             sub_state = ?7,
                             exit_code = ?8,
                             main_pid = ?9,
                             memory_peak = ?10,
                             pids_peak = ?11,
                             stdout = ?12,
                             stderr = ?13
                         WHERE jobid = ?14 AND taskid = ?15")?;
                    stmt.execute(params![
                        unitname,
                        start_time,
                        end_time,
                        mono_start_time,
                        mono_end_time,
                        active_state,
                        sub_state,
                        exit_code,
                        main_pid,
                        memory_peak,
                        pids_peak,
                        stdout,
                        stderr,
                        jobid,
                        taskid
                    ])?;
                }
            }
            else {
                stmt = conn.prepare(
                    "INSERT INTO status_task
                     (
                       jobid,
                       taskid,
                       unitname,
                       start_time,
                       end_time,
                       mono_start_time,
                       mono_end_time,
                       active_state,
                       sub_state,
                       exit_code,
                       main_pid,
                       memory_peak,
                       pids_peak,
                       stdout,
                       stderr
                     )
                     VALUES(
                       ?1,
                       ?2,
                       ?3,
                       ?4,
                       ?5,
                       ?6,
                       ?7,
                       ?8,
                       ?9,
                       ?10,
                       ?11,
                       ?12,
                       ?13,
                       ?14,
                       ?15
                     )")?;
                stmt.execute(params![
                    jobid,
                    taskid,
                    unitname,
                    start_time,
                    end_time,
                    mono_start_time,
                    mono_end_time,
                    active_state,
                    sub_state,
                    exit_code,
                    main_pid,
                    memory_peak,
                    pids_peak,
                    stdout,
                    stderr
                ])?;
            }

            Ok(vec![DbResult::InsertSuccess])
        },
        _ => Err(rusqlite::Error::InvalidQuery),
    }
}



/// Main function to call
/// to get or insert
/// data from or to database.
pub async fn execute(pool: &Pool, q: DbQuery) -> Result<Vec<DbResult>, Error> {
    let pool = pool.clone();

    let conn = web::block(move || pool.get())
        .await?
        .map_err(error::ErrorInternalServerError)?;

    web::block(move || {
        match q {
            DbQuery::GetL3qdValidate => get_l3qd_validate(conn),
            DbQuery::GetNodeL3qdValidate => get_node_l3qd_validate(conn),
            DbQuery::GetAllTasks => get_all_tasks(conn),
            DbQuery::InsertTask{..} => insert_task(conn, q),
            _ => Err(rusqlite::Error::InvalidQuery),
        }
    })
        .await?
        .map_err(error::ErrorInternalServerError)
}


/// To be used with status thread
/// as the thread is not async
pub fn execute_status(pool: &Pool, q: DbQuery) -> Result<Vec<DbResult>, rusqlite::Error> {
    let pool = pool.clone();

    let conn = match pool.get() {
        Ok(c) => c,
        Err(_) => return
            Err(rusqlite::Error::InvalidQuery),
    };

    match q {
        DbQuery::GetL3qdValidate => get_l3qd_validate(conn),
        DbQuery::GetNodeL3qdValidate => get_node_l3qd_validate(conn),
        DbQuery::GetAllTasks => get_all_tasks(conn),
        DbQuery::DeleteAllTasks => delete_all_tasks(conn),
        DbQuery::InsertTask{..} => insert_task(conn, q),
        _ => Err(rusqlite::Error::InvalidQuery),
    }
}


/// Initializes database
/// and creates tables
/// if they do not exist.
/// Returns a pool to be used
/// for queries.
pub fn init_db(db_path: String) -> Result<Pool, String> {
    let manager = r2d2_sqlite::SqliteConnectionManager::file(db_path);
    let pool;
    match Pool::new(manager) {
	Ok(p) => pool = p,
	Err(e) => {
	    return Err(
		format!("Failed to create database connection pool: {e}")
	    );
	}
    }

    let conn;
    match pool.get() {
	Ok(c) => conn = c,
	Err(e) => {
	    return Err(
		format!("Failed to get connection to database: {e}")
	    );
	}
    }

    match conn.execute(
	"CREATE TABLE IF NOT EXISTS validate
         (
             id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,
             ip TEXT NOT NULL,
             name TEXT NOT NULL,
             key TEXT NOT NULL
         );", params![]) {
        Ok(_) => (),
        Err(e) => {
	    return Err(
		format!("Failed to create table validate_l3qd: {e}")
	    );
	}
    }

    match conn.execute(
	"CREATE TABLE IF NOT EXISTS status_task
         (
             id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,
             jobid INT,
             taskid INT,
             unitname TEXT,
             start_time TEXT,
             end_time TEXT,
             mono_start_time INT,
             mono_end_time INT,
             active_state TEXT,
             sub_state TEXT,
             exit_code INT,
             main_pid INT,
             memory_peak INT,
             pids_peak INT,
             stdout TEXT,
             stderr TEXT
         );", params![]) {
        Ok(_) => (),
        Err(e) => {
	    return Err(
		format!("Failed to create table status_task: {e}")
	    );
	}
    }
    
    Ok(pool)
}
